package com.ruoyi.ldzlsc.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * 任务状态WebSocket处理器
 * 
 * @author ldzl
 * @date 2025-01-23
 */
@Component
public class TaskStatusWebSocketHandler implements WebSocketHandler {

    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    // 存储所有WebSocket连接
    private static final CopyOnWriteArraySet<WebSocketSession> sessions = new CopyOnWriteArraySet<>();
    
    // 存储用户ID与WebSocket会话的映射
    private static final ConcurrentHashMap<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        
        // 从查询参数中获取用户ID
        String userId = getUserIdFromSession(session);
        if (userId != null) {
            userSessions.put(userId, session);
            System.out.println("WebSocket连接建立，用户ID: " + userId + ", 会话ID: " + session.getId());
        } else {
            System.out.println("WebSocket连接建立，会话ID: " + session.getId());
        }
        
        // 发送连接成功消息
        sendMessage(session, new WebSocketMessage<>("connection", "连接成功", null));
    }

    @Override
    public void handleMessage(WebSocketSession session, org.springframework.web.socket.WebSocketMessage<?> message) throws Exception {
        if (message instanceof TextMessage) {
            TextMessage textMessage = (TextMessage) message;
            try {
                WebSocketMessage<?> wsMessage = objectMapper.readValue(textMessage.getPayload(), WebSocketMessage.class);
                System.out.println("收到WebSocket消息: " + wsMessage);

                // 处理心跳消息
                if ("ping".equals(wsMessage.getType())) {
                    sendMessage(session, new WebSocketMessage<>("pong", "心跳响应", null));
                }
            } catch (Exception e) {
                System.err.println("解析WebSocket消息失败: " + e.getMessage());
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("WebSocket传输错误，会话ID: " + session.getId() + ", 错误: " + exception.getMessage());
        removeSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        System.out.println("WebSocket连接关闭，会话ID: " + session.getId() + ", 状态: " + closeStatus);
        removeSession(session);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 移除会话
     */
    private void removeSession(WebSocketSession session) {
        sessions.remove(session);
        
        // 从用户会话映射中移除
        String userIdToRemove = null;
        for (String userId : userSessions.keySet()) {
            if (userSessions.get(userId).equals(session)) {
                userIdToRemove = userId;
                break;
            }
        }
        if (userIdToRemove != null) {
            userSessions.remove(userIdToRemove);
        }
    }

    /**
     * 从会话中获取用户ID
     */
    private String getUserIdFromSession(WebSocketSession session) {
        try {
            String query = session.getUri().getQuery();
            if (query != null && query.contains("userId=")) {
                String[] params = query.split("&");
                for (String param : params) {
                    if (param.startsWith("userId=")) {
                        return param.substring(7);
                    }
                }
            }
        } catch (Exception e) {
            System.err.println("获取用户ID失败: " + e.getMessage());
        }
        return null;
    }

    /**
     * 发送消息给指定会话
     */
    public void sendMessage(WebSocketSession session, WebSocketMessage<?> message) {
        try {
            if (session.isOpen()) {
                String jsonMessage = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(jsonMessage));
            }
        } catch (IOException e) {
            System.err.println("发送WebSocket消息失败: " + e.getMessage());
        }
    }

    /**
     * 广播消息给所有连接
     */
    public void broadcastMessage(WebSocketMessage<?> message) {
        String jsonMessage;
        try {
            jsonMessage = objectMapper.writeValueAsString(message);
        } catch (Exception e) {
            System.err.println("序列化WebSocket消息失败: " + e.getMessage());
            return;
        }

        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(jsonMessage));
                }
            } catch (IOException e) {
                System.err.println("广播WebSocket消息失败: " + e.getMessage());
                removeSession(session);
            }
        }
    }

    /**
     * 发送消息给指定用户
     */
    public void sendMessageToUser(String userId, WebSocketMessage<?> message) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            sendMessage(session, message);
        } else {
            System.out.println("用户 " + userId + " 的WebSocket连接不存在或已关闭");
        }
    }

    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return sessions.size();
    }
}
